#------run.py
import torch
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from models import TransformerGenerator # 确保models.py在同一目录下

class FoggyImageGenerator:
    def __init__(self, model_path, device='cuda'):
        """
        初始化加雾生成器
        :param model_path: 预训练生成器模型路径
        :param device: 推理设备 (cuda/cpu)
        """
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
        self.model = self._load_model(model_path).to(self.device)
        self.transform = transforms.Compose([
            transforms.ToTensor(), # 自动将PIL图像归一化到[0,1]
        ])

    def _load_model(self, model_path):
        """加载预训练生成器模型"""
        model = TransformerGenerator(in_channels=3, out_channels=3, nhead=8)
        state_dict = torch.load(model_path, map_location=self.device)
        model.load_state_dict(state_dict)
        model.eval()
        return model

    def _preprocess(self, image):
        """图像预处理"""

        # 确保图像为RGB格式
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # 转换为Tensor并添加批次维度
        tensor = self.transform(image).unsqueeze(0).to(self.device)
        return tensor

    def _postprocess(self, tensor):
        """结果后处理"""

        # 移除非批次维度并转换为numpy数组
        output = tensor.squeeze().cpu().detach().numpy()

        # 调整通道顺序 (C x H x W) -> (H x W x C)
        output = np.transpose(output, (1, 2, 0))

        # 转换为0-255范围
        output = (output * 255).astype(np.uint8)
        return Image.fromarray(output)

    def add_fog(self, input_image):
        """
        为输入图像添加雾霾效果
        :param input_image: PIL.Image对象 或 图像路径
        :return: 加雾后的PIL.Image对象
        """

        # 加载图像（如果输入的是路径）
        if isinstance(input_image, str):
            input_image = Image.open(input_image)

        # 记录原始尺寸
        original_size = input_image.size

        # 预处理
        input_tensor = self._preprocess(input_image)

        # 模型推理
        with torch.no_grad():
            output_tensor = self.model(input_tensor)

        # 后处理
        output_image = self._postprocess(output_tensor)

        # 保持原始尺寸（必要时进行缩放）
        if output_image.size != original_size:
            output_image = output_image.resize(original_size, Image.BILINEAR)

        return output_image

if __name__ == "__main__":
    import argparse

    # 命令行参数解析
    parser = argparse.ArgumentParser(description='为输入图像添加雾霾效果')
    parser.add_argument('--input', type=str, required=True, help='输入图像路径')
    parser.add_argument('--output', type=str, required=True, help='输出图像路径')
    parser.add_argument('--model', type=str, default='checkpoints_true/generator_epoch_200.pth',
                        help='预训练生成器模型路径')
    parser.add_argument('--device', type=str, choices=['cuda', 'cpu'], default='cuda',
                        help='推理设备 (默认: cuda)')
    args = parser.parse_args()

    # 初始化生成器
    fog_generator = FoggyImageGenerator(args.model, device=args.device)

    # 处理图像
    result = fog_generator.add_fog(args.input)

    # 保存结果
    result.save(args.output)
    print(f"处理完成，结果已保存至 {args.output}")

